"""
OCR服务模块 - 基于RapidOCR API的文字识别服务
集成了OCR接入指南中的核心功能
"""

import requests
import datetime
import json
import logging
import os
import subprocess
import time
from PIL import ImageGrab, Image
from typing import Optional, Dict, Any, Tuple, List


class OCRService:
    """OCR服务封装类"""
    
    def __init__(self, api_url: str = "http://127.0.0.1:9003/ocr"):
        self.api_url = api_url
        self.base_url = api_url.replace('/ocr', '')
        self.setup_logging()
        
    def setup_logging(self):
        """配置日志系统"""
        # 创建logs目录
        log_dir = os.path.join(os.path.dirname(__file__), '..', 'logs')
        os.makedirs(log_dir, exist_ok=True)
        
        # 配置日志
        log_file = os.path.join(log_dir, 'ocr.log')
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.StreamHandler(),  # 控制台输出
                logging.FileHandler(log_file, mode='w', encoding='utf-8')  # 文件输出(覆盖写入)
            ]
        )
        self.logger = logging.getLogger(__name__)
        self.logger.info("OCR服务日志系统初始化完成")
    
    def check_service_status(self) -> bool:
        """检查OCR服务状态"""
        try:
            response = requests.get(f"{self.base_url}/docs", timeout=10)
            if response.status_code == 200:
                self.logger.info("✅ OCR服务运行正常")
                return True
        except requests.exceptions.ConnectionError:
            self.logger.warning("⚠️ OCR服务连接失败")
        except Exception as e:
            self.logger.error(f"❌ 服务状态检查异常: {e}")
        return False
    
    def ensure_service_running(self) -> bool:
        """确保OCR服务正常运行"""
        # 先检查是否已经运行
        if self.check_service_status():
            return True
            
        self.logger.info("正在尝试启动OCR服务...")
        
        try:
            # 尝试后台启动服务
            subprocess.Popen(['rapidocr_api'], 
                           stdout=subprocess.DEVNULL, 
                           stderr=subprocess.DEVNULL)
            
            # 等待服务启动(最多等20秒)
            for i in range(10):
                time.sleep(2)
                if self.check_service_status():
                    self.logger.info("✅ OCR服务启动成功")
                    return True
                self.logger.info(f"等待服务启动... ({i+1}/10)")
            
            self.logger.error("❌ OCR服务启动超时")
            return False
            
        except FileNotFoundError:
            self.logger.error("❌ 未找到rapidocr_api命令，请检查是否已安装")
            return False
        except Exception as e:
            self.logger.error(f"❌ 启动服务异常: {e}")
            return False
    
    def take_screenshot(self, save_path: str = None, bbox: Tuple[int, int, int, int] = None, enhance: bool = True) -> Optional[str]:
        """
        截图功能（支持图片增强处理）
        Args:
            save_path: 保存路径，为空则自动生成
            bbox: 截图区域 (left, top, right, bottom)，为空则全屏截图
            enhance: 是否进行图片增强处理（放大和清晰化）
        Returns:
            截图文件路径
        """
        try:
            # 生成保存路径
            if save_path is None:
                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                img_dir = os.path.join(os.path.dirname(__file__), '..', 'images', 'screenshots')
                os.makedirs(img_dir, exist_ok=True)
                save_path = os.path.join(img_dir, f"screenshot_{timestamp}.png")
            
            # 截图
            if bbox:
                screenshot = ImageGrab.grab(bbox=bbox)
                self.logger.info(f"区域截图完成: {bbox}")
            else:
                screenshot = ImageGrab.grab()
                self.logger.info("全屏截图完成")
            
            # 图片增强处理
            if enhance:
                screenshot = self.enhance_image(screenshot)
                self.logger.info("图片增强处理完成")
            
            # 保存图片
            screenshot.save(save_path)
            self.logger.info(f"截图已保存: {save_path}")
            return save_path
            
        except Exception as e:
            self.logger.error(f"截图失败: {e}")
            return None
    
    def enhance_image(self, image: Image.Image, scale_factor: float = 3.0) -> Image.Image:
        """
        图片增强处理：放大和清晰化
        Args:
            image: PIL图片对象
            scale_factor: 放大倍数，默认3倍
        Returns:
            增强后的图片对象
        """
        try:
            from PIL import ImageFilter, ImageEnhance
            
            # 1. 放大图片（使用高质量重采样）
            original_size = image.size
            new_size = (int(original_size[0] * scale_factor), int(original_size[1] * scale_factor))
            enhanced_img = image.resize(new_size, Image.LANCZOS)
            
            # 2. 转换为灰度图（提高OCR准确性）
            enhanced_img = enhanced_img.convert('L')
            
            # 3. 增强对比度
            enhancer = ImageEnhance.Contrast(enhanced_img)
            enhanced_img = enhancer.enhance(1.5)  # 增强1.5倍对比度
            
            # 4. 增强锐度
            enhancer = ImageEnhance.Sharpness(enhanced_img)
            enhanced_img = enhancer.enhance(2.0)  # 增强2倍锐度
            
            # 5. 去噪处理
            enhanced_img = enhanced_img.filter(ImageFilter.MedianFilter(size=3))
            
            self.logger.info(f"图片增强：{original_size} -> {new_size}, 对比度+50%, 锐度+100%")
            return enhanced_img
            
        except Exception as e:
            self.logger.warning(f"图片增强失败，使用原图: {e}")
            return image
    
    def call_ocr_api(self, img_path: str) -> Optional[Dict[str, Any]]:
        """
        调用OCR API识别文字
        Args:
            img_path: 图片文件路径
        Returns:
            识别结果字典
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(img_path):
                self.logger.error(f"图片文件不存在: {img_path}")
                return None
            
            self.logger.info(f"开始OCR识别: {img_path}")
            
            # 调用API（使用正确的参数名 'image_file'）
            with open(img_path, 'rb') as f:
                files = {'image_file': (os.path.basename(img_path), f, 'image/png')}
                response = requests.post(self.api_url, files=files, timeout=30)
            
            self.logger.info(f"API响应状态码: {response.status_code}")
            
            if response.status_code == 200:
                result = response.json()
                text_count = len(result) if result else 0
                self.logger.info(f"✅ OCR识别成功，识别到 {text_count} 个文本区域")
                
                # 保存识别结果
                self.save_ocr_result(result, img_path)
                self.save_ocr_result_json(result, img_path)  # 同时保存JSON格式
                return result
            else:
                self.logger.error(f"❌ API调用失败: {response.status_code}")
                self.logger.error(f"错误详情: {response.text}")
                return None
                
        except Exception as e:
            self.logger.error(f"OCR调用异常: {e}")
            return None
    
    def save_ocr_result(self, result: Dict[str, Any], img_path: str):
        """保存OCR识别结果到文本文件"""
        try:
            # 生成结果文件路径
            base_name = os.path.splitext(os.path.basename(img_path))[0]
            result_dir = os.path.join(os.path.dirname(__file__), '..', 'logs')
            txt_path = os.path.join(result_dir, f"{base_name}_ocr_result.txt")
            
            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write("=" * 60 + "\n")
                f.write("OCR识别结果报告\n")
                f.write("=" * 60 + "\n")
                f.write(f"时间戳: {datetime.datetime.now()}\n")
                f.write(f"图片路径: {img_path}\n")
                f.write(f"识别区域总数: {len(result) if result else 0}\n")
                f.write("=" * 60 + "\n\n")
                
                # 写入详细结果
                if result:
                    for key, value in result.items():
                        text = value.get('rec_txt', '')
                        score = value.get('score', 0)
                        bbox = value.get('dt_boxes', [])
                        

                else:
                    f.write("未识别到任何文本\n")
            
            self.logger.info(f"识别结果已保存: {txt_path}")
            
        except Exception as e:
            self.logger.error(f"保存识别结果失败: {e}")
    
    def save_ocr_result_json(self, result: Dict[str, Any], img_path: str):
        """保存OCR识别结果到JSON文件"""
        try:
            # 生成JSON文件路径
            base_name = os.path.splitext(os.path.basename(img_path))[0]
            result_dir = os.path.join(os.path.dirname(__file__), '..', 'ocr_results')
            os.makedirs(result_dir, exist_ok=True)
            json_path = os.path.join(result_dir, f"{base_name}_ocr_result.json")
            
            # 构建完整的结果对象
            full_result = {
                'timestamp': datetime.datetime.now().isoformat(),
                'image_path': img_path,
                'total_regions': len(result) if result else 0,
                'ocr_data': result or {},
                'enhanced_data': self.extract_structured_data(result) if result else {}
            }
            
            # 保存为JSON文件（格式化输出，方便查看）
            with open(json_path, 'w', encoding='utf-8') as f:
                json.dump(full_result, f, ensure_ascii=False, indent=2)
            
            self.logger.info(f"📄 JSON结果已保存: {json_path}")
            
            # 清理旧文件，只保留最近10张
            self._cleanup_old_files()
            
        except Exception as e:
            self.logger.error(f"保存JSON结果失败: {e}")
    
    def _cleanup_old_files(self):
        """清理旧的OCR结果和截图文件，只保留最近10张"""
        try:
            # 清理OCR结果文件
            result_dir = os.path.join(os.path.dirname(__file__), '..', 'ocr_results')
            if os.path.exists(result_dir):
                json_files = [f for f in os.listdir(result_dir) if f.endswith('_ocr_result.json')]
                json_files.sort(key=lambda x: os.path.getctime(os.path.join(result_dir, x)), reverse=True)
                
                # 删除超过10个的文件
                for file_to_delete in json_files[10:]:
                    file_path = os.path.join(result_dir, file_to_delete)
                    os.remove(file_path)
                    self.logger.debug(f"删除旧OCR结果: {file_to_delete}")
            
            # 清理截图文件
            img_dir = os.path.join(os.path.dirname(__file__), '..', 'images', 'screenshots')
            if os.path.exists(img_dir):
                img_files = [f for f in os.listdir(img_dir) if f.startswith('screenshot_') and f.endswith('.png')]
                img_files.sort(key=lambda x: os.path.getctime(os.path.join(img_dir, x)), reverse=True)
                
                # 删除超过10个的文件
                for file_to_delete in img_files[10:]:
                    file_path = os.path.join(img_dir, file_to_delete)
                    os.remove(file_path)
                    self.logger.debug(f"删除旧截图: {file_to_delete}")
                    
            self.logger.debug("文件清理完成")
            
        except Exception as e:
            self.logger.error(f"文件清理失败: {e}")
    
    def extract_structured_data(self, result: Dict[str, Any]) -> Dict[str, Any]:
        """从OCR结果提取结构化数据"""
        try:
            structured_data = {
                'text_items': [],
                'potential_table_data': []
            }
            
            if not result:
                return structured_data
            
            # 提取所有文本项并按位置排序
            text_items = []
            for key, value in result.items():
                text = value.get('rec_txt', '').strip()
                bbox = value.get('dt_boxes', [])
                score = value.get('score', 0)
                
                if text and bbox and len(bbox) >= 4:
                    # 计算中心坐标
                    y_coords = [point[1] for point in bbox]
                    x_coords = [point[0] for point in bbox]
                    avg_y = sum(y_coords) / len(y_coords)
                    avg_x = sum(x_coords) / len(x_coords)
                    
                    text_items.append({
                        'id': key,
                        'text': text,
                        'confidence': score,
                        'position': {'x': avg_x, 'y': avg_y},
                        'bbox': bbox
                    })
            
            # 按位置排序（先按Y坐标，再按X坐标）
            text_items.sort(key=lambda x: (x['position']['y'], x['position']['x']))
            structured_data['text_items'] = text_items
            
            # 尝试提取表格数据
            structured_data['potential_table_data'] = self.extract_table_structure(text_items)
            
            return structured_data
            
        except Exception as e:
            self.logger.error(f"提取结构化数据失败: {e}")
            return {'text_items': [], 'potential_table_data': []}
    
    def extract_table_structure(self, text_items: List[Dict]) -> List[Dict]:
        """尝试从文本项中提取表格结构"""
        try:
            if not text_items:
                return []
            
            # 按行分组（Y坐标相近的为同一行）
            rows = []
            current_row = []
            current_y = text_items[0]['position']['y']
            y_threshold = 20  # 减小Y坐标差异阈值，提高行分组精度
            
            self.logger.info(f"🔍 开始行分组，Y阈值={y_threshold}，总文本数={len(text_items)}")
            
            for item in text_items:
                y_diff = abs(item['position']['y'] - current_y)
                if y_diff <= y_threshold:
                    current_row.append(item)
                    self.logger.debug(f"  添加到当前行: '{item['text']}' (Y差异: {y_diff:.1f})")
                else:
                    if current_row:
                        # 按X坐标排序当前行
                        current_row.sort(key=lambda x: x['position']['x'])
                        rows.append(current_row)
                        self.logger.info(f"完成行{len(rows)}: {len(current_row)}个元素")
                    current_row = [item]
                    current_y = item['position']['y']
                    self.logger.info(f"开始新行: '{item['text']}' (Y={current_y:.1f}, Y差异: {y_diff:.1f})")
            
            # 添加最后一行
            if current_row:
                current_row.sort(key=lambda x: x['position']['x'])
                rows.append(current_row)
            
            # 转换为表格数据
            table_data = []
            for row_idx, row in enumerate(rows):
                row_data = {
                    'row_index': row_idx,
                    'columns': []
                }
                for col_idx, item in enumerate(row):
                    text = item['text']
                    # 股票代码字段修正：第二列（column_index=1）为股票代码，全部转为大写
                    if col_idx == 1:  # 第二列是股票代码
                        # 将常见的OCR错误字符修正为大写字母
                        text = self.fix_stock_code(text)
                    
                    row_data['columns'].append({
                        'column_index': col_idx,
                        'text': text,
                        'confidence': item['confidence']
                    })
                table_data.append(row_data)
                
                # 添加调试日志
                row_texts = [col['text'] for col in row_data['columns']]
                self.logger.info(f"🔍 表格行 {row_idx + 1}: {len(row_texts)}列 - {row_texts}")
                
                # 详细的位置调试信息
                for col_idx, item in enumerate(row):
                    self.logger.info(f"  列{col_idx + 1}: '{item['text']}' 位置(x={item['position']['x']:.1f}, y={item['position']['y']:.1f})")
            
            self.logger.info(f"表格提取完成：共 {len(table_data)} 行数据")
            return table_data
            
        except Exception as e:
            self.logger.error(f"提取表格结构失败: {e}")
            return []
    
    def find_text_in_result(self, result: Dict[str, Any], keyword: str, exact_match: bool = False) -> List[Dict]:
        """
        在识别结果中查找指定关键词
        Args:
            result: OCR识别结果
            keyword: 要查找的关键词
            exact_match: 是否精确匹配
        Returns:
            包含匹配信息的列表
        """
        matches = []
        
        if not result:
            return matches
            
        for key, value in result.items():
            text = value.get('rec_txt', '')
            
            # 检查是否匹配
            if exact_match:
                is_match = keyword == text
            else:
                is_match = keyword in text
                
            if is_match:
                bbox = value.get('dt_boxes', [])
                score = value.get('score', 0)
                
                # 计算中心点坐标
                center_x, center_y = self.calculate_center_point(bbox)
                
                match_info = {
                    'key': key,
                    'text': text,
                    'score': score,
                    'bbox': bbox,
                    'center': (center_x, center_y)
                }
                matches.append(match_info)
                
        
        if not matches:
            self.logger.warning(f"未找到关键词: {keyword}")
            
        return matches
    
    def calculate_center_point(self, bbox: List[List[int]]) -> Tuple[int, int]:
        """计算边界框的中心点"""
        if not bbox or len(bbox) < 4:
            return (0, 0)
            
        try:
            x_coords = [point[0] for point in bbox]
            y_coords = [point[1] for point in bbox]
            center_x = int(sum(x_coords) / len(x_coords))
            center_y = int(sum(y_coords) / len(y_coords))
            return (center_x, center_y)
        except:
            return (0, 0)
    
    def fix_stock_code(self, text: str) -> str:
        """
        修正股票代码中的OCR识别错误
        股票代码只包含大写字母，不包含数字和小写字母
        """
        if not text:
            return text
            
        original_text = text
        fixed_text = text
        
        # 小写字母修正规则：基于字形相似度的智能修正
        lowercase_corrections = {
            'l': 'I',   # 小写l最容易被误认为I（如MIAX被识别成MlAX）
            'i': 'I',   # 小写i应该是大写I
            'o': 'O',   # 小写o应该是大写O
            'c': 'C',   # 小写c应该是大写C
            's': 'S',   # 小写s应该是大写S
            'x': 'X',   # 小写x应该是大写X
            'z': 'Z',   # 小写z应该是大写Z
            'v': 'V',   # 小写v应该是大写V
            'w': 'W',   # 小写w应该是大写W
            'u': 'U',   # 小写u应该是大写U
            'p': 'P',   # 小写p应该是大写P
            'k': 'K',   # 小写k应该是大写K
            'f': 'F',   # 小写f应该是大写F
            't': 'T',   # 小写t应该是大写T
            'y': 'Y',   # 小写y应该是大写Y
            'm': 'M',   # 小写m应该是大写M
            'n': 'N',   # 小写n应该是大写N
            'h': 'H',   # 小写h应该是大写H
            'b': 'B',   # 小写b应该是大写B
            'd': 'D',   # 小写d应该是大写D
            'g': 'G',   # 小写g应该是大写G
            'j': 'J',   # 小写j应该是大写J
            'q': 'Q',   # 小写q应该是大写Q
            'r': 'R',   # 小写r应该是大写R
            'a': 'A',   # 小写a应该是大写A
            'e': 'E',   # 小写e应该是大写E
        }
        
        # 应用小写字母修正规则
        for wrong, correct in lowercase_corrections.items():
            fixed_text = fixed_text.replace(wrong, correct)
        
        # 数字修正规则：数字在股票代码中不应该存在
        number_corrections = {
            '0': 'O',  # 数字0应该是字母O
            '1': 'I',  # 数字1应该是字母I  
            '2': 'Z',  # 数字2可能是字母Z
            '5': 'S',  # 数字5可能是字母S
            '8': 'B',  # 数字8可能是字母B
        }
        
        # 应用数字修正规则
        for wrong, correct in number_corrections.items():
            fixed_text = fixed_text.replace(wrong, correct)
        
        # 过滤掉所有非字母字符（股票代码只应该包含字母）
        import re
        fixed_text = re.sub(r'[^A-Z]', '', fixed_text)
        
        if fixed_text != original_text:
            self.logger.info(f"📝 股票代码修正: '{original_text}' -> '{fixed_text}' (小写字母和数字已修正)")
            
        return fixed_text

    def ocr_area(self, bbox: Tuple[int, int, int, int]) -> Optional[Dict[str, Any]]:
        """
        对指定区域进行OCR识别
        Args:
            bbox: 区域坐标 (left, top, right, bottom)
        Returns:
            识别结果
        """
        try:
            # 先截图指定区域（开启图片增强）
            img_path = self.take_screenshot(bbox=bbox, enhance=True)
            if not img_path:
                return None
            
            # 调用OCR识别
            result = self.call_ocr_api(img_path)
            return result
            
        except Exception as e:
            self.logger.error(f"区域OCR识别失败: {e}")
            return None
